{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4f237efb",
   "metadata": {},
   "source": [
    "## 昇思+昇腾开发板：软硬结合玩转大模型实践能力认证（中级）\n",
    "\n",
    "**环境准备：**\n",
    "\n",
    "开发者拿到香橙派开发板后，首先需要进行硬件资源确认、镜像烧录以及CANN和MindSpore版本的升级，才可运行该案例，具体如下：\n",
    "\n",
    "|**香橙派AIpro**|**镜像**|**CANN Toolkit/Kernels**|**MindSpore**|**MindSpore NLP**|\n",
    "|:-------:|:-------:|:-------:|:-------:|:-------:|\n",
    "|20T 24G|Ubuntu|8.0.0beta1|2.5.0|0.4分支|\n",
    "\n",
    "- CANN检查与升级：参考[链接](https://www.mindspore.cn/tutorials/zh-CN/r2.6.0/orange_pi/environment_setup.html#3-cann%E5%8D%87%E7%BA%A7)\n",
    "- MindSpore检查与升级：参考[链接](https://www.mindspore.cn/tutorials/zh-CN/r2.6.0/orange_pi/environment_setup.html#4-mindspore%E5%8D%87%E7%BA%A7)\n",
    "- MindSpore NLP安装命令：\n",
    "    ```bash\n",
    "    pip install git+https://github.com/mindspore-lab/mindnlp.git@0.4\n",
    "    ```\n",
    "\n",
    "**场景说明：** 在本任务中，我们将实现DeepSeek-R1-Distill-Qwen-1.5B的解码工程，模型在输出logits后，进行包含贪心搜索、top-p采样、temperature采样等选的文本解码策略，最终根据提供的prompt进行文本生成。\n",
    "\n",
    "**任务目标：** 本任务旨在考核对解码算法的掌握，以及基于MindSpore的开发实践。当前notebook文件已完成了代码工程框架的搭建，并在4个关键模块进行了代码扣空，开发者在昇思官网(https://www.mindspore.cn) 中找到**MindSpore2.5.0版本的接口文档，重点关注mindspore.mint接口文档**，参考文档补全空缺处的代码，保证执行脚本全流程跑通，最终根据开发者自行设计的prompt输出模型生成内容。\n",
    "\n",
    "**参考文档：**\n",
    "\n",
    "开发者可在昇思官网(https://www.mindspore.cn) 中找到**MindSpore2.5.0版本的接口文档，重点关注mindspore.mint接口文档**\n",
    "\n",
    "> 空缺处可通过多行代码实现，不局限于一行"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cde296eb",
   "metadata": {},
   "source": [
    "### 模型实例化\n",
    "\n",
    "模型链接：https://modelers.cn/models/MindSpore-Lab/DeepSeek-R1-Distill-Qwen-1.5B-FP16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "943c3e5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入包\n",
    "import mindspore\n",
    "from mindnlp.transformers import AutoTokenizer, AutoModelForCausalLM, StaticCache\n",
    "from mindspore import mint\n",
    "import numpy as np\n",
    "\n",
    "# 开启同步，用于定位问题，调试完毕后建议关闭同步\n",
    "# mindspore.set_context(pynative_synchronize=True)\n",
    "\n",
    "# 模型实例化\n",
    "model_id = \"MindSpore-Lab/DeepSeek-R1-Distill-Qwen-1.5B-FP16\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id, mirror=\"modelers\")\n",
    "model = AutoModelForCausalLM.from_pretrained(model_id, ms_dtype=mindspore.float16, low_cpu_mem_usage=True, mirror=\"modelers\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b733610e",
   "metadata": {},
   "source": [
    "### **考点1：Top-p 采样实现**\n",
    "Top-p 采样是一种根据累积概率动态选择候选词集的解码策略，通过设定阈值p保留概率最高的部分词进行采样，以平衡生成问题的多样性和相关性。\n",
    "\n",
    "**实现逻辑如下：**<br>\n",
    "**1.概率排序：** 将模型输出的词汇概率按从高到低排序<br>\n",
    "**2.累积计算：** 对排序后的概率进行累加，生成累积概率分布<br>\n",
    "**3.阈值截断：** 保留累积概率首次超过设定值p的最小词集合<br>\n",
    "**4.子集采样：** 仅从阶段后的词集中依概率随机选取下一个词<br>\n",
    "\n",
    "**要求：** 请根据上述top-p逻辑，和代码中注释的提示，完成top-p的逻辑实现。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aea05b89",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Top-p采样实现\n",
    "def sample_top_p(probs, p=0.9):\n",
    "    \"\"\"\n",
    "    Top-p采样函数，用于生成文本时选择下一个token。\n",
    "    \"\"\"\n",
    "    # 按概率降序排序\n",
    "    probs_np = probs.asnumpy()\n",
    "    sorted_indices = np.argsort(-probs_np, axis=-1)\n",
    "    sorted_probs = np.take_along_axis(probs_np, sorted_indices, axis=-1)\n",
    "    # 累积计算\n",
    "    cumulative_probs = np.cumsum(sorted_probs, axis=-1)\n",
    "     # 构建掩码，阈值截断\n",
    "    mask = cumulative_probs - sorted_probs > p\n",
    "    sorted_probs[mask] = 0.0\n",
    "    sorted_probs /= np.sum(sorted_probs, axis=-1, keepdims=True)\n",
    "    # 子集采样\n",
    "    # >>>>>>> 填空：请根据提示，开发top-p采样代码 <<<<<<<\n",
    "    # 注意香橙派上支持的数据类型为float16和int32\n",
    "    sorted_probs = ________ # 转换为mindspore.Tensor\n",
    "    sorted_indices = ________ # 转换为mindspore.Tensor\n",
    "    new_token_idx = ________ # 获取采样得到的索引，请使用mint.multinomial接口\n",
    "    # 获取token\n",
    "    new_token = mindspore.ops.gather(\n",
    "        sorted_indices, new_token_idx, \n",
    "        axis=-1, batch_dims=1\n",
    "    )\n",
    "    return new_token"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "587da211",
   "metadata": {},
   "source": [
    "### **考点2： 单token解码函数实现**\n",
    "在解码过程中，我们首先获取模型单token的前向计算logits，其次根据解码策略（温度采样、top-p采样，或是贪心搜索），获取下一个词。\n",
    "我们将如上过程写作一个函数decode_one_token。\n",
    "\n",
    "**要求：** 如下代码已实现了单token的前向计算，请根据如下要求，补齐代码，实现对输出logits的采样和检索。\n",
    "1. 可自由控制是否执行温度采样，及采样程度，可和top-p采样配合使用\n",
    "2. 可自由控制是否执行Top-p采样，及采样程度，可和温度采样配合使用\n",
    "3. 如果温度采样和Top-p采样均不执行，默认执行贪心搜索\n",
    "\n",
    "**温度采样调节：** 直接使用模型输出的logits进行采样往往会导致概率分布过于尖锐、生成结果缺乏多样性等问题，为此引入温度调节机制，对应公式如下：\n",
    "\n",
    "$$logits = logits / temperature$$\n",
    "\n",
    "当温度参数大于1时，会压缩原始logits差异、平缓概率分布、增加生成的多样性，例如：\n",
    "$$temperature = 1.5$$\n",
    "$$logits = [5.0, 3.0, 2.0]$$\n",
    "$$logits_{new} = [3.33, 2.0, 1.33]$$\n",
    "\n",
    "当温度参数小于1时，会放大原始logits差异、锐化概率分布、增加生成的确定性，例如：\n",
    "$$temperature = 0.7$$\n",
    "$$logits = [5.0, 3.0, 2.0]$$\n",
    "$$logits_{new} = [7.14, 4.29, 2.86]$$\n",
    "\n",
    "我们鼓励开发者使用不同的temperature进行实验。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d618cb67",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 单token解码函数实现\n",
    "def decode_one_tokens_logits(\n",
    "        model, cur_token, input_pos, \n",
    "        cache_position, past_key_values\n",
    "):\n",
    "    \"\"\"单个token的解码函数，返回logits\"\"\"\n",
    "    logits = model(\n",
    "        cur_token,\n",
    "        position_ids=input_pos,\n",
    "        cache_position=cache_position,\n",
    "        past_key_values=past_key_values,\n",
    "        return_dict=False,\n",
    "        use_cache=True\n",
    "    )[0]\n",
    "    return logits\n",
    "\n",
    "\n",
    "def decode_one_tokens(\n",
    "        model, cur_token, input_pos, \n",
    "        cache_position, past_key_values, \n",
    "        temperature, top_p\n",
    "):\n",
    "    \"\"\"单个token的解码函数，由logits、温度和Top_p选择合适的token\"\"\"\n",
    "    # 模型前向计算结果\n",
    "    logits = decode_one_tokens_logits(\n",
    "        model, cur_token, input_pos,\n",
    "        cache_position, past_key_values\n",
    "    )\n",
    "\n",
    "    # >>>>>>> 填空：补齐单token解码逻辑 <<<<<<<\n",
    "    # 若使用温度采样 + Top-p采样，请先使用温度采样方法 调整模型输出的logits，\n",
    "    # 将调整后的logits转化为概率分布后，使用sample_top_p接口得到new_token\n",
    "    # 若使用贪心搜索，请直接返回模型输出的logits对应概率最高的token即可      \n",
    "    new_token = ________  # 实现可以为多行代码\n",
    "    return new_token"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5521546",
   "metadata": {},
   "source": [
    "### **考点3： prefill阶段实现**\n",
    "在prefill阶段，模型在首词处理输入序列时，会预先计算并缓存当前所有位置的key、value矩阵，对后续的自回归生成提供加速基础。\n",
    "\n",
    "prefill阶段逻辑如下：<br>\n",
    "**1. 静态缓存**：通过创建静态缓存实例化past_key_values<br>\n",
    "**2. 创建generated_ids：** 创建generated_ids，此为最终生成的内容，后续在decode阶段会不断进行更新<br>\n",
    "**3. 前向计算：** 初始前向计算获取首个logits<br>\n",
    "**4. next token生成：** 生成第一个新token<br>\n",
    "**5. 更新generated_ids：** 将新生成的token更新入generated_ids中<br>\n",
    "\n",
    "**要求：** 如下代码已完成了静态缓存的实现，请参考上述prefill阶段的逻辑，完成prefill的代码实现。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f3d473c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# >>>>>>> 填空：请设计一个（或多个）prompt输入 <<<<<<<\n",
    "prompts = [________]\n",
    "\n",
    "# tokenization\n",
    "inputs = tokenizer(prompts, return_tensors=\"ms\", padding=True)\n",
    "\n",
    "# 生成参数配置\n",
    "# >>>>>>> 填空：请根据提示，设计合适的参数配置 <<<<<<<\n",
    "NUM_TOKENS_TO_GENERATE = ________   # 每个输入要生成的token数量\n",
    "TEMPERATURE = ________  # 温度参数（控制生成多样性）\n",
    "TOP_P = ________  # Top-p采样阈值\n",
    "\n",
    "\n",
    "batch_size, seq_length = inputs[\"input_ids\"].shape\n",
    "\n",
    "# 静态缓存：创建静态缓存（用于加速自回归生成）\n",
    "past_key_values = StaticCache(\n",
    "    config=model.config, max_batch_size=len(prompts), max_cache_len=512, dtype=model.dtype\n",
    ")\n",
    "\n",
    "# 生成的第一个token的位置\n",
    "cache_position = mint.arange(seq_length, dtype=mindspore.int32)\n",
    "\n",
    "# >>>>>>> 填空：补齐prefill逻辑 <<<<<<<\n",
    "# 创建generated_ids，用于存放最终生成的内容，后续会不断更新\n",
    "# generated_ids是一个存放token_id的容器，长度为输入序列长度 + 期望生成的token数量 + 1个生成结束的标识符\n",
    "# generated_ids存放的数据应为：输入序列 + 生成序列 + 结束符，数据类型为int32，每个数据对应词表中的一个token，\n",
    "# 后续通过解码器将模型生成的token编码序列，解码回自然语言\n",
    "generated_ids = ________  # 实现可以为多行代码\n",
    "\n",
    "# 初始前向传播获取首个logits\n",
    "logits = model(\n",
    "    **inputs, \n",
    "    cache_position=cache_position, \n",
    "    past_key_values=past_key_values,\n",
    "    return_dict=False, \n",
    "    use_cache=True\n",
    ")[0]\n",
    "\n",
    "# 生成第一个新token，参考前面生成new_token的逻辑\n",
    "________  # 实现可以为多行代码\n",
    "\n",
    "# 更新generated_ids\n",
    "________"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee83c50b",
   "metadata": {},
   "source": [
    "### **考点4：decode阶段实现**\n",
    "在decode阶段，模型在自回归生成每个新的token时，复用之前缓存的KV并仅计算当前步的注意力，从而降低推理复杂度。\n",
    "\n",
    "decode阶段逻辑如下：<br>\n",
    "**1. 构建循环：** 构建自回归循环<br>\n",
    "**2. next token预测：** 生成当前新的token<br>\n",
    "**3. 更新cache position**<br>\n",
    "\n",
    "Qwen模型推理机制属于step-by-step，每次仅会生成一个新token，且每个token的生成都需要前文所有的token信息，表示公式如下：\n",
    "$$token_{n+1} = decode\\_one\\_tokens(token_{1...n})$$\n",
    "\n",
    "**要求：** 请参考上述逻辑补齐代码，完成decode阶段的代码实现。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ba55e7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 自回归生成循环\n",
    "cache_position = mindspore.tensor([seq_length + 1])\n",
    "# >>>>>>> 填空：补齐自回归循环 <<<<<<<\n",
    "# 请调用decode_one_tokens接口获取下一个token，并在generated_ids容器中完成对应位置token_id的更新\n",
    "for _ in range(1, NUM_TOKENS_TO_GENERATE):\n",
    "    ________  # 实现可以为多行代码\n",
    "\n",
    "# 将最终生成的tokens（generated_ids）转换为文本\n",
    "text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)\n",
    "print(text)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mindspore",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
